# encoding: utf-8

import torch
import onnxruntime as ort


class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = torch.nn.Linear(10, 10)
        self.fc = torch.nn.Linear(10, 1)
        self.dropout = torch.nn.Dropout(0.4)
        self.activate = torch.nn.SELU()

    def forward(self, x):
        x = self.linear1(x)
        x = self.activate(x)
        x = self.dropout(x)
        x = self.fc(x)
        return x


x = torch.randn(2, 10)

print(x)

model = Model()

with torch.no_grad():

    torch.onnx.export(model, x, "model2.onnx", opset_version=9, input_names=["input"], output_names=["output"])
